#!/usr/bin/env python
# coding=utf-8
# __author__ = 'Yunchao Ling'

import os


def parse_vcf(filepath: str):
    var_list = []
    infile = open(filepath, "r")
    for line in infile:
        line = line.rstrip()
        if not line.startswith("#"):
            splitline = line.split("\t")
            pos = int(splitline[1])
            if pos >= 21563 and pos <= 25384:
                variant = "%s:%s:%s" % (splitline[1], splitline[3], splitline[4])
                var_list.append(variant)
    infile.close()
    return var_list


def get_s_var(folder: str, out_path: str):
    var_ht = {}
    for lst in os.listdir(folder):
        if os.path.isfile(os.path.join(folder, lst)) and lst.endswith("vcf"):
            var_list = parse_vcf(os.path.join(folder, lst))
            if var_list:
                virus_id = lst[:-4]
                for variant in var_list:
                    if variant in var_ht:
                        var_ht[variant].append(virus_id)
                    else:
                        var_ht[variant] = [virus_id]
    outfile = open(out_path, "w")
    outfile.write("pos\tref\talt\tvirus_ids\n")
    for variant in var_ht:
        split_variant = variant.split(":")
        outfile.write(
            "%s\t%s\t%s\t%s\n" % (split_variant[0], split_variant[1], split_variant[2], ",".join(var_ht[variant])))
        outfile.flush()
    outfile.close()


if __name__ == '__main__':
    get_s_var("/Users/genesis/Downloads/non-human_annoation_vcf", "/Users/genesis/Downloads/non_human_var_S.tsv")
